import sys
import torch
from transformers import AutoModelForCausalLM
import pdb
def compare_models(model_path1, model_path2):
    try:
        model1 = AutoModelForCausalLM.from_pretrained(model_path1, trust_remote_code=True)
        model2 = AutoModelForCausalLM.from_pretrained(model_path2, trust_remote_code=True)
    except:
        from transformers import Qwen2AudioForConditionalGeneration
        model1 = Qwen2AudioForConditionalGeneration.from_pretrained(model_path1, trust_remote_code=True)
        model2 = Qwen2AudioForConditionalGeneration.from_pretrained(model_path2, trust_remote_code=True)

    pdb.set_trace()
    model1_params = model1.state_dict()
    model2_params = model2.state_dict()

    if model1_params.keys() != model2_params.keys():
        print("The models have different architectures.")
        return

    for key in model1_params:
        diff_sum = torch.sum(torch.abs(model1_params[key] - model2_params[key]))
        diff_avg = diff_sum / model1_params[key].numel()
        print(f"{key} diff_sum: {diff_sum.item()}")
        print(f"{key} diff_avg: {diff_avg.item()}")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python compare_models.py <model_path1> <model_path2>")
    else:
        compare_models(sys.argv[1], sys.argv[2])
